from scipy.signal import find_peaks
import pandas as pd
import matplotlib.pyplot as plt
import mplfinance as mpf
# import warnings
# warnings.filterwarnings("ignore", category=UserWarning, module="mplfinance")
def merge_number(pd):
        # 提取每个元素的第一个值
    prices = [item[0] for item in pd]

    # 初始化合并结果列表
    merged_prices = []
    i = 0
    while i < len(prices):
        current_price = prices[i]
        generation = 1
        merged = [current_price]
        j = i + 1
        while j < len(prices):
            next_price = prices[j]
            # 计算差值百分比
            diff_percentage = abs((next_price - current_price) / current_price)
            if diff_percentage <= 0.02:
                merged.append(next_price)
                generation += 1
                j += 1
            else:
                break
        # 将合并结果和代数添加到最终结果列表
        merged_prices.append((merged, generation))
        i = j
    return merged_prices
    # # 输出合并结果
    # for prices, gen in merged_prices:
    #     print(f"第 {gen} 代: {prices}")

def calculate_strong_peaks(df):
    strong_peaks, _ = find_peaks(df['high'],distance=20)
    strong_peaks_values = df.iloc[strong_peaks][['high', 'date']].values.tolist()
    # 包括52周最高价作为额外的强峰值
    # yearly_high = df['high'].iloc[:20].max()
    # strong_peaks_values.append([yearly_high, None])
    print("强峰值（阻力位）:", strong_peaks_values)
    return strong_peaks_values

def calculate_troughs_values(df):
    # 找到'低'价格数据中的强谷值
    strong_troughs, _ = find_peaks(-df['low'],distance=20)
    # 提取强谷值的对应低值
    strong_troughs_values = df.iloc[strong_troughs][['low', 'date']].values.tolist()
    # 包括52周最低价作为额外的强谷值
    # yearly_low = df['low'].iloc[:20].min()
    # strong_troughs_values.append([yearly_low, None])
    # print("强谷值（支撑位）:", strong_troughs_values)
    
    return strong_troughs_values

def plot_kline_with_support_resistance(df, strong_peaks_values, strong_troughs_values):
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)

    # 准备K线图数据
    ohlc = df[['open', 'high', 'low', 'close']]

    # 准备支撑位和压力位的数据（严格保持日期-价格对应）
    resistance_data = [(x[0], pd.Timestamp(x[1])) for x in strong_peaks_values if x[1] is not None]
    resistance_prices = [x[0] for x in resistance_data]
    resistance_dates = [x[1] for x in resistance_data]
    
    support_data = [(x[0], pd.Timestamp(x[1])) for x in strong_troughs_values if x[1] is not None]
    support_prices = [x[0] for x in support_data]
    support_dates = [x[1] for x in support_data]

    # 创建与主图索引匹配的Series（关键修复）
    resistance_series = pd.Series(resistance_prices, index=resistance_dates).reindex(ohlc.index, method='nearest')
    support_series = pd.Series(support_prices, index=support_dates).reindex(ohlc.index, method='nearest')
    # 绘制K线图和成交量柱状图，设置上涨和下跌颜色
    mc = mpf.make_marketcolors(up='r', down='g')  # 设置上涨颜色为红色，下跌颜色为绿色
    s = mpf.make_mpf_style(marketcolors=mc)

    # 创建K线图（修改参数为正确的形式）
    ap = [
        mpf.make_addplot(resistance_series, type='scatter', marker='^', color='r'),
        mpf.make_addplot(support_series, type='scatter', marker='v', color='g')
    ]
    mpf.plot(ohlc, type='candle', addplot=ap, title='K line graph', ylabel='price',datetime_format='%Y-%m-%d',style=s)

if __name__ == "__main__":
    import os
    file_path = os.path.join('csv', '113631_full.csv')
    csv_data = pd.read_csv(file_path)[:120]
    strong_peaks = calculate_strong_peaks(csv_data)
    strong_troughs = calculate_troughs_values(csv_data)
    print(merge_number(strong_troughs))
    plot_kline_with_support_resistance(csv_data, strong_peaks, strong_troughs)